from dataclasses import dataclass
from typing import Optional, List, Dict, Any

@dataclass
class ScenarioMetadata:
    """
    故障场景元数据
    """
    scenario_id: str               # 场景唯一ID
    folder_name: str               # 原始文件夹名称
    anomaly_component: str         # 故障组件
    anomaly_reason: str            # 故障类型
    dataset_type: str              # "train" 或 "test"
    trace_count: int               # 可用Trace数量
    level_folder: Optional[str] = None  # 输出文件夹名称，使用scenario_id


@dataclass
class TraceSample:
    """
    Trace样本记录
    """
    trace_id: str                  # Trace唯一标识
    scenario_id: str               # 所属场景ID
    length: int                    # Trace长度
    sample_type: str               # "train" 或 "test"
    level_folder: str              # 输出文件夹名称，使用scenario_id


class DatasetMetadata:
    """
    数据集元数据
    """
    def __init__(self, config: Dict[str, Any]):
        self.total_scenarios = 0
        self.train_scenarios = 0
        self.test_scenarios = 0
        self.total_samples = 0
        self.train_samples = 0
        self.test_samples = 0
        self.config = config
        
    def update_from_scenarios(self, scenarios: List[ScenarioMetadata]):
        """从场景列表更新统计数据"""
        self.total_scenarios = len(scenarios)
        self.train_scenarios = sum(1 for s in scenarios if s.dataset_type == "train")
        self.test_scenarios = sum(1 for s in scenarios if s.dataset_type == "test")
        
    def update_from_samples(self, samples: List[TraceSample]):
        """从样本列表更新统计数据"""
        self.total_samples = len(samples)
        self.train_samples = sum(1 for s in samples if s.sample_type == "train")
        self.test_samples = sum(1 for s in samples if s.sample_type == "test")
        
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "total_scenarios": self.total_scenarios,
            "train_scenarios": self.train_scenarios,
            "test_scenarios": self.test_scenarios,
            "total_samples": self.total_samples,
            "train_samples": self.train_samples,
            "test_samples": self.test_samples,
            "config": self.config
        }